
//go:build ignore

#include <linux/bpf.h>
#include <bpf/bpf_helpers.h>
#include <linux/if_ether.h>
//#include <linux/ip.h> // 包含用于定义 struct iphdr 的头文件
//#include <arpa/inet.h>
#include <bpf/bpf_endian.h>
#include <linux/in.h>


//#define TARGET_IP 0x0a020102

// 10.2.1.1
//#define TARGET_IP 0x0A020101
struct iphdr {
  __u8        ihl:4,
        version:4;
  __u8        tos;
  __be16  tot_len;
  __be16  id;
  __be16  frag_off;
  __u8        ttl;
  __u8        protocol;
  __sum16 check;
  __be32  saddr;
  __be32  daddr;
};

struct {
    __uint(type, BPF_MAP_TYPE_ARRAY); 
    __type(key, __u32);
    __type(value, __u64);
    __uint(max_entries, 100);
} pkt_count SEC(".maps");
struct {
    __uint(type, BPF_MAP_TYPE_ARRAY);
    __type(key, __u32);
    __type(value, __u32);
    __uint(max_entries, 1);
} ip_blacklist SEC(".maps");

// count_packets atomically increases a packet counter on every invocation.
SEC("xdp") 
int count_packets(struct xdp_md *ctx)  {

   __u32 key    = 0;
   __u64 *count = bpf_map_lookup_elem(&pkt_count, &key);
   if (count) {
      __sync_fetch_and_add(count, 1);
   }

    __u32 key1    = 0;
    __u32 *count1 = bpf_map_lookup_elem(&ip_blacklist, &key1);
    //struct xdp_md *ctx;
    void *data_end = (void *)(long)ctx->data_end;
    void *data = (void *)(long)ctx->data;

    struct ethhdr *eth = data;
    struct iphdr *ip;
    if((data + sizeof(*eth)) > data_end) {
        // 错误：缺少以太网头
        return XDP_DROP;
    }

    if (eth->h_proto == bpf_htons(ETH_P_IP)) {
        ip = (struct iphdr *)(data + sizeof(*eth));
        if ((data + sizeof(*eth) + sizeof(*ip)) > data_end) {
          // 错误：缺少 IP 头
          return XDP_DROP;
        }
        __be32 src_ip = ip->saddr;
        if (count1 && *count1 == src_ip) { // IPv4 地址 如果源IP 是 10.2.1.1
                 return XDP_DROP;
        }
    }
      return XDP_PASS;

  }

char __license[] SEC("license") = "Dual MIT/GPL";
